// =================================================================================
// 模块: 必要的头文件和全局常量定义
// Module: Necessary header files and global constant definitions
// =================================================================================
#include <stdio.h>
#include <string.h>
#include <stdlib.h> 
#include <math.h> 

// --- 全局常量定义 ---
#define IMG_ROWS 28
#define IMG_COLS 28
#define FLATTEN_SIZE (IMG_ROWS * IMG_COLS) // 784
#define HIDDEN1_NEURONS 128
#define OUTPUT_NEURONS 10

// --- 文件路径定义 ---
#define W1_FILE_PATH "input_files/W1.txt"
#define B1_FILE_PATH "input_files/B1.txt"
#define W2_FILE_PATH "input_files/W2.txt"
#define B2_FILE_PATH "input_files/B2.txt"
#define EXAMPLE_IMAGE_DIR "example/"


// =================================================================================
// 模块: PGM图像文件读取 和 矩阵/向量 转换 (Phase 1, Functions 1, 2, 3, 4)
// 来源于 "pgm读取（phase1中的1和2）" 和 "偏值读取（phase1中的3和4）" 提供的代码。
// =================================================================================

// 1. 读取PGM文件到矩阵
void pgm_to_matrix(const char *file_path, const int n_rows, const int n_cols,
                   double A[n_rows][n_cols]) {
    FILE *fp = fopen(file_path, "rb");
    if (fp == NULL) {
        printf("Error! Image file '%s' not found or cannot be opened.\n", file_path); 
        for(int i=0; i<n_rows; ++i) for(int j=0; j<n_cols; ++j) A[i][j] = 0.0;
        return;
    }

    // 跳过PGM文件头（4行）
    char buffer[256];
    for (int i = 0; i < 4; i++) {
        if (!fgets(buffer, sizeof(buffer), fp)) {
            fprintf(stderr, "Invalid PGM header in file: %s\n", file_path);
            fclose(fp);
            for(int r=0; r<n_rows; ++r) for(int c=0; c<n_cols; ++c) A[r][c] = 0.0;
            return;
        }
    }

    // 读取像素数据
    for (int i = 0; i < n_rows; i++) {
        for (int j = 0; j < n_cols; j++) {
            int pixel = fgetc(fp);
            if (pixel == EOF) {
                fprintf(stderr, "Unexpected end of file in PGM: %s (at row %d, col %d)\n", file_path, i, j);
                fclose(fp);
                A[i][j] = 0.0; 
                return;
            }
            A[i][j] = (double)pixel / 255.0;
        }
    }
    fclose(fp);
}

// 2. 矩阵展平为向量
void matrix_to_vector(const int n_rows, const int n_cols,
                      const double matrix[n_rows][n_cols],
                      double vector[1][n_rows * n_cols]) {
    int idx = 0;
    for (int i = 0; i < n_rows; i++) {
        for (int j = 0; j < n_cols; j++) {
            vector[0][idx++] = matrix[i][j];
        }
    }
}

// 3. 向量重塑为矩阵
void vector_to_matrix(const int n_rows, const int n_cols,
                      const double vector[1][n_rows * n_cols],
                      double matrix[n_rows][n_cols]) {
    if (n_rows <= 0 || n_cols <= 0) {
        printf("Error: Invalid matrix dimensions for vector_to_matrix.\n"); 
        return;
    }
    int idx = 0; // "pgm读取" 版本使用 idx，"偏值读取" 版本使用 vector_index 计算，这里统一为 idx
    for (int i = 0; i < n_rows; i++) {
        for (int j = 0; j < n_cols; j++) {
            matrix[i][j] = vector[0][idx++];
        }
    }
}

// 4. 读取文本数据到向量
void read_text_data(const char *filename, const int n, double array[1][n]) {
    // 打开文件
    FILE *fp = fopen(filename, "r");
    if (fp == NULL) {
        fprintf(stderr, "Error: Cannot open file %s for reading text data.\n", filename); // 英文错误信息
        for(int i=0; i<n; ++i) array[0][i] = 0.0; // 安全初始化
        return;
    }

    // 读取文件中的数据
    for (int i = 0; i < n; i++) {
        if (fscanf(fp, "%lf", &array[0][i]) != 1) {
            fprintf(stderr, "Error: Failed to read data from file %s (at element %d).\n", filename, i); // 英文错误信息
            for (int k=i; k<n; ++k) array[0][k] = 0.0; // 将剩余部分清零
            fclose(fp);
            return;
        }
    }
    // 关闭文件
    fclose(fp);
}


// =================================================================================
// 模块: 矩阵运算 (Phase 1, Functions 5, 6)
// 来源于 "pgm读取（phase1中的1和2）" 和 "矩阵乘法和加法（phase1中的5和6）" 提供的代码。
// =================================================================================

// 5. 矩阵加法
// (来自 "矩阵乘法和加法" 部分)
void add_matrices(const int n_rows, const int n_cols,
                  const double mat1[n_rows][n_cols],
                  const double mat2[n_rows][n_cols],
                  double result[n_rows][n_cols]) {
    for (int i = 0; i < n_rows; i++) {
        for (int j = 0; j < n_cols; j++) {
            result[i][j] = mat1[i][j] + mat2[i][j];
        }
    }
}

// 6. 矩阵乘法
void multiply_matrices(const int numRowsMat1, const int numColsMat2, const int commonDimension,
                       const double matrix1[][commonDimension],
                       const double matrix2[][numColsMat2],
                       double productMatrix[][numColsMat2]) {
    for (int i = 0; i < numRowsMat1; i++) {
        for (int j = 0; j < numColsMat2; j++) {
            productMatrix[i][j] = 0.0;
        }
    }

    // 执行矩阵乘法.
    for (int i = 0; i < numRowsMat1; i++) {
        for (int j = 0; j < numColsMat2; j++) {
            for (int k = 0; k < commonDimension; k++) {
                productMatrix[i][j] += matrix1[i][k] * matrix2[k][j];
            }
        }
    }
}


// =================================================================================
// 模块: 激活函数 (Phase 1, Functions 7, 8)
// Module: Activation Functions (Phase 1, Functions 7, 8)
// =================================================================================

// 7. ReLU 激活函数
void activation_function_relu(const int SIZE, const double F[1][SIZE], double G[1][SIZE]) {
    for (int m = 0; m < SIZE; m++) {
        // 取 F[m] 和 0 的最大值
        G[0][m] = fmax(F[0][m], 0.0);
    }
}

// 8. SoftMax 函数
void softmax_function(const int n_cols, const double l[1][n_cols], double s[1][n_cols]) {
    if (n_cols <= 0) {
        fprintf(stderr, "Error: Invalid dimension for Softmax input.\n");
        return;
    }
    // 找到最大值
    double max_val = l[0][0];
    for (int i = 1; i < n_cols; i++) {
        if (l[0][i] > max_val) max_val = l[0][i];
    }

    // 计算指数和
    double sum_exp = 0.0; 
    for (int i = 0; i < n_cols; i++) {
        s[0][i] = exp(l[0][i] - max_val); 
        sum_exp += s[0][i];
    }

    // 归一化
    if (sum_exp == 0) { // 防止除以零
        fprintf(stderr, "Warning: Softmax sum_exp is zero. Outputting uniform distribution.\n");
        for (int i = 0; i < n_cols; i++) {
            s[0][i] = 1.0 / n_cols;
        }
        return;
    }
    for (int i = 0; i < n_cols; i++) {
        s[0][i] /= sum_exp;
    }
}

// =================================================================================
// 模块: 神经网络层逻辑 (部分来源于 "第一隐藏层权重计算..." 代码)
// Module: Neural Network Layer Logic (partially from "第一隐藏层权重计算..." code)
// =================================================================================


// =================================================================================
// 模块: 核心识别流程 和 主函数
// Module: Core Recognition Pipeline and Main Function
// =================================================================================

/**
 * 核心识别函数: 单个手写数字识别 (阶段二逻辑)
 */
int recognize_single_digit(const char *pgmFilePath) {
    // --- 阶段二 步骤中使用的矩阵 ---

    // 步骤 1: 读取PGM图像并将其重塑为 1x784 向量B
    double matrixA[IMG_ROWS][IMG_COLS];       // 28x28 图像矩阵 (28x28 image matrix)
    pgm_to_matrix(pgmFilePath, IMG_ROWS, IMG_COLS, matrixA);

    double vectorB[1][FLATTEN_SIZE];          // 1x784 (matrixA reshaped)
    matrix_to_vector(IMG_ROWS, IMG_COLS, matrixA, vectorB);

    // --- 第一隐藏层 ---

    // 步骤 2: 应用第一隐藏层权重 W1 (文件中的 C 矩阵)
    // W1.txt 包含 784*128 = 100352 个元素, 重塑为 matrixC_W1 (784x128)
    // matrixD = vectorB * matrixC_W1 (1x784 * 784x128 = 1x128)
    double W1_flat[1][FLATTEN_SIZE * HIDDEN1_NEURONS];
    read_text_data(W1_FILE_PATH, FLATTEN_SIZE * HIDDEN1_NEURONS, W1_flat);

    double matrixC_W1[FLATTEN_SIZE][HIDDEN1_NEURONS]; // 对应项目文档的 C 矩阵（第一层权重）
    vector_to_matrix(FLATTEN_SIZE, HIDDEN1_NEURONS, W1_flat, matrixC_W1);


    double matrixD[1][HIDDEN1_NEURONS]; // 对应项目文档的 D 矩阵
    multiply_matrices(1, HIDDEN1_NEURONS, FLATTEN_SIZE, vectorB, matrixC_W1, matrixD);

    // 步骤 3: 应用第一隐藏层偏置 B1 (文件中的 E 向量)
    // matrixF = matrixD + vectorE_B1 (1x128 + 1x128 = 1x128)
    double vectorE_B1[1][HIDDEN1_NEURONS]; // B1.txt (1x128), 对应项目文档的 E 矩阵 (偏置)
    read_text_data(B1_FILE_PATH, HIDDEN1_NEURONS, vectorE_B1);

    double matrixF[1][HIDDEN1_NEURONS]; // 对应项目文档的 F 矩阵
    add_matrices(1, HIDDEN1_NEURONS, matrixD, vectorE_B1, matrixF);

    // 步骤 4: 激活 (ReLU)
    // vectorG = ReLU(matrixF) (1x128)
    double vectorG[1][HIDDEN1_NEURONS]; // 对应项目文档的 G 矩阵
    activation_function_relu(HIDDEN1_NEURONS, matrixF, vectorG);


    // --- 第二隐藏层 ---
    // Second Hidden Layer

    // 步骤 5: 应用第二隐藏层权重 W2
    // W2.txt 包含 128*10 = 1280 个元素, 重塑为 W2_matrix (128x10)
    // matrixH = vectorG * W2_matrix (1x128 * 128x10 = 1x10)
    double W2_flat[1][HIDDEN1_NEURONS * OUTPUT_NEURONS];
    read_text_data(W2_FILE_PATH, HIDDEN1_NEURONS * OUTPUT_NEURONS, W2_flat);

    double W2_matrix[HIDDEN1_NEURONS][OUTPUT_NEURONS];
    vector_to_matrix(HIDDEN1_NEURONS, OUTPUT_NEURONS, W2_flat, W2_matrix);

    double matrixH[1][OUTPUT_NEURONS]; // 对应项目文档的 H 矩阵
    multiply_matrices(1, OUTPUT_NEURONS, HIDDEN1_NEURONS, vectorG, W2_matrix, matrixH);

    // 步骤 6: 应用第二隐藏层偏置 B2
    // matrixL = matrixH + B2_bias_vector (1x10 + 1x10 = 1x10)
    double B2_bias_vector[1][OUTPUT_NEURONS]; // B2.txt (1x10)
    read_text_data(B2_FILE_PATH, OUTPUT_NEURONS, B2_bias_vector);

    double matrixL[1][OUTPUT_NEURONS]; // 对应项目文档的 L 矩阵
    add_matrices(1, OUTPUT_NEURONS, matrixH, B2_bias_vector, matrixL);

    // 步骤 7: SoftMax
    // vectorS = SoftMax(matrixL) (1x10)
    double vectorS[1][OUTPUT_NEURONS]; // 对应项目文档的 S 矩阵
    softmax_function(OUTPUT_NEURONS, matrixL, vectorS);

    // 步骤 8: 找出识别的数字 (S中具有最高概率的索引)
    int recognizedDigit = 0;
    double maxProbability = vectorS[0][0];
    for (int i = 1; i < OUTPUT_NEURONS; i++) {
        if (vectorS[0][i] > maxProbability) {
            maxProbability = vectorS[0][i];
            recognizedDigit = i;
        }
    }
    return recognizedDigit;
}


/**
 * 主函数: 用户交互和流程控制
 * Main function: User interaction and flow control
 */
int main() {
    char userChoice = 'Y';
    char filenamesBuffer[1024]; 
    char *currentFilenameToken;
    char *trimmedToken; // 用于存储去除前后空格后的token

    printf("Group Project - C Handwritten Digit Recognition\n"); // 项目标题 (Project Title)
    printf("=========================================================\n");

    while (userChoice == 'Y' || userChoice == 'y') {
        printf("\nEnter image filename(s) (e.g., p1.pgm or p1.pgm,p2.pgm,p3.pgm): "); // 提示用户输入 (Prompt user for input)
        if (fgets(filenamesBuffer, sizeof(filenamesBuffer), stdin) == NULL) {
            fprintf(stderr, "Error reading user input!\n");
            break; 
        }

        filenamesBuffer[strcspn(filenamesBuffer, "\n")] = '\0'; // 移除换行符

        if (strlen(filenamesBuffer) == 0) {
            printf("No filename entered.\n");
            goto ask_to_continue;
        }

        currentFilenameToken = strtok(filenamesBuffer, ","); // 主要以逗号分隔

        while (currentFilenameToken != NULL) {
            trimmedToken = currentFilenameToken;
            while (*trimmedToken == ' ' || *trimmedToken == '\t') {
                trimmedToken++;
            }
            char *end = trimmedToken + strlen(trimmedToken) - 1;
            while (end >= trimmedToken && (*end == ' ' || *end == '\t')) {
                *end-- = '\0';
            }

            if (strlen(trimmedToken) == 0) { 
                currentFilenameToken = strtok(NULL, ","); 
                continue;
            }

            char fullImagePath[512];
            snprintf(fullImagePath, sizeof(fullImagePath), "%s%s", EXAMPLE_IMAGE_DIR, trimmedToken);

            FILE *checkFile = fopen(fullImagePath, "rb");
            if (checkFile == NULL) {
                 printf("invalid file.\n"); 
            } else {
                fclose(checkFile); 
                printf("Processing %s ... ", trimmedToken);
                fflush(stdout);
                int digit = recognize_single_digit(fullImagePath);

                printf("%s: %d\n", trimmedToken, digit);
            }

            currentFilenameToken = strtok(NULL, ","); 
        }

    ask_to_continue:
        printf("\nDo you want to continue recognition? (Y/N): ");
        char tempChoiceBuffer[10];
        if (fgets(tempChoiceBuffer, sizeof(tempChoiceBuffer), stdin) != NULL) {
            tempChoiceBuffer[strcspn(tempChoiceBuffer, "\n")] = 0; 
            if (strlen(tempChoiceBuffer) > 0) {
                userChoice = tempChoiceBuffer[0];
            } else {
                userChoice = 'N'; 
            }
        } else {
            userChoice = 'N'; 
        }
        // 转换为大写
        if (userChoice >= 'a' && userChoice <= 'z') {
            userChoice -= ('a' - 'A'); 
        }
        if (userChoice != 'Y') { // 只有明确输入 Y 才继续，其他都视为 N
            userChoice = 'N';
        }
    }

    printf("\nProgram exited. Thank you for using!\n"); 
    return 0;
}